use actix::prelude::*;
use futures::future::join_all;
use std::collections::HashMap;
use std::sync::Arc;

// Define messages
struct ExecuteTask {
    input: usize,
    env: Arc<EnvVar>,
}

impl Message for ExecuteTask {
    type Result = Result<usize, ()>;
}

// Define the task actor
struct TaskActor {
    value: usize,
    name: String,
}

impl TaskActor {
    fn new(value: usize, name: &str) -> Self {
        TaskActor {
            value,
            name: name.to_string(),
        }
    }
}

impl Actor for TaskActor {
    type Context = Context<Self>;
}

impl Handler<ExecuteTask> for TaskActor {
    type Result = Result<usize, ()>;

    fn handle(&mut self, msg: ExecuteTask, _: &mut Context<Self>) -> Self::Result {
        let base = msg.env.get::<usize>("base").unwrap();
        let result = self.value * base + msg.input;
        Ok(result)
    }
}

// Define the environment variable structure
struct EnvVar {
    vars: HashMap<String, usize>,
}

impl EnvVar {
    fn new() -> Self {
        EnvVar {
            vars: HashMap::new(),
        }
    }

    fn set(&mut self, key: &str, value: usize) {
        self.vars.insert(key.to_string(), value);
    }

    fn get<T: std::str::FromStr>(&self, key: &str) -> Option<T> {
        self.vars.get(key).and_then(|v| v.to_string().parse().ok())
    }
}

// Define a structure to hold the DAG configuration
struct DagConfig {
    tasks: HashMap<String, (usize, Vec<String>)>,
}

impl DagConfig {
    fn new() -> Self {
        let mut tasks = HashMap::new();
        tasks.insert("A".to_string(), (1, vec![]));
        tasks.insert("B".to_string(), (2, vec!["A".to_string()]));
        tasks.insert("C".to_string(), (4, vec!["A".to_string()]));
        tasks.insert("D".to_string(), (8, vec!["A".to_string()]));
        tasks.insert("E".to_string(), (16, vec!["B".to_string(), "C".to_string()]));
        tasks.insert("F".to_string(), (32, vec!["C".to_string(), "D".to_string()]));
        tasks.insert("G".to_string(), (64, vec!["B".to_string(), "E".to_string(), "F".to_string()]));
        DagConfig { tasks }
    }
}

// Define the scheduler actor
struct SchedulerActor {
    actors: HashMap<String, Addr<TaskActor>>,
    config: DagConfig,
    env: Arc<EnvVar>,
}

impl SchedulerActor {
    fn new(actors: HashMap<String, Addr<TaskActor>>, config: DagConfig, env: Arc<EnvVar>) -> Self {
        SchedulerActor { actors, config, env }
    }
}

impl Actor for SchedulerActor {
    type Context = Context<Self>;
}

impl Handler<ExecuteTask> for SchedulerActor {
    type Result = ResponseActFuture<Self, Result<usize, ()>>;

    fn handle(&mut self, msg: ExecuteTask, _: &mut Context<Self>) -> Self::Result {
        let futures = self.config.tasks.keys().map(|name| {
            let actor = self.actors.get(name).unwrap().clone();
            actor.send(ExecuteTask {
                input: msg.input,
                env: self.env.clone(),
            })
        });

        Box::pin(
            join_all(futures)
                .into_actor(self)
                .map(|results, _actor, _ctx| {
                    let mut total = 0;
                    for res in results {
                        if let Ok(Ok(val)) = res {
                            total += val;
                        }
                    }
                    Ok(total)
                }),
        )
    }
}

#[actix::main]
async fn main() {
    // Create environment
    let mut env = EnvVar::new();
    env.set("base", 2);

    // Load DAG configuration
    let config = DagConfig::new();

    // Create a map to hold task actors
    let mut actors = HashMap::new();

    // Initialize actors based on configuration
    for (name, (value, _)) in &config.tasks {
        let actor = TaskActor::new(*value, name).start();
        actors.insert(name.clone(), actor);
    }

    // Create and start the scheduler actor
    let env = Arc::new(env);
    let scheduler = SchedulerActor::new(actors, config, env.clone()).start();

    // Dispatch tasks
    match scheduler.send(ExecuteTask { input: 0, env }).await {
        Ok(result) => println!("The result is {}.", result.unwrap()),
        Err(_) => println!("Failed to execute DAG."),
    }
}
